import json
import os

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.datasets import make_blobs

from src.util import util


class Participant:

    def __init__(self, idx, pos):
        self.idx = idx
        self.pos = pos

    def calc_distance(self, other):
        return np.linalg.norm(self.pos - other.pos).astype(int)


class Task:

    def __init__(self, idx, size, cycles_per_bit, weight_time, weight_energy):
        self.idx = idx
        self.size, self.cycles_per_bit = size, cycles_per_bit
        self.weight_time, self.weight_energy = weight_time, weight_energy

    @classmethod
    def factory(cls, idx, min_size, max_size):
        size = np.random.randint(min_size, max_size) << 3   # 任务大小，单位 bit
        cycles_per_bit = np.random.randint(1, 3) * 100      # 每 bit 所需的时钟周期数目
        weight_time = np.random.uniform(1e-3, 1)  # 时间权重
        weight_energy = 1 - weight_time                     # 能耗权重
        return Task(idx, size, cycles_per_bit, weight_time, weight_energy)

    def __str__(self):
        return f"Task:{self.idx}, size:{self.size}, cycles_per_bit:{self.cycles_per_bit}"


class UE(Participant):

    def __init__(self, idx, pos, task, cpu_frequency, trans_power, coefficient_performance):
        super().__init__(idx, pos)
        self.task = task                    # 当前任务信息
        self.cpu_frequency = cpu_frequency  # 本地计算的最大 CPU 频率
        self.trans_power = trans_power      # 任务卸载的传输功率
        self.coefficient_performance = coefficient_performance  # 能效系数
        self.part = -1          # 所属 Part ID
        self.offloading = 0     # 是否选择卸载
        self.time_total = 0     # 处理当前任务所需总时间，考虑卸载决策的取值
        self.cost = 0           # 处理当前任务的总成本，考虑卸载决策的取值
        self.t_local = 0        # 当前任务在本地计算所需的时间
        self.t_wait_drone = 0   # 当前 UD 等待 UAV 到达的时间
        self.t_off = 0          # 当前任务任务卸载所需的时间
        self.t_exe = 0          # 当前任务卸载给 UAV 后等待执行结束所需的时间

    def set_task(self, task):
        self.task = task
        self.t_local = (task.size * task.cycles_per_bit) / self.cpu_frequency

    def calc_cost(self, t_wait_drone, trans_rate, drone_cpu_frequency):
        if self.offloading:
            t_off = self.task.size / trans_rate
            t_exe = (self.task.size * self.task.cycles_per_bit) / drone_cpu_frequency
            time_total = t_wait_drone + t_off + t_exe
        else:
            time_total = self.t_local
        return time_total

    def calc_exe_time(self, t_wait_drone, trans_rate, drone_cpu_frequency):
        if not self.offloading:
            return self.t_local
        t_off = self.task.size / trans_rate
        t_exe = (self.task.size * self.task.cycles_per_bit) / drone_cpu_frequency
        return t_wait_drone + t_off + t_exe

    def execute_action(self, t_wait_drone, trans_rate, drone_cpu_frequency):
        if self.offloading:
            self.t_wait_drone = t_wait_drone
            self.t_off = self.task.size / trans_rate
            self.t_exe = (self.task.size * self.task.cycles_per_bit) / drone_cpu_frequency
            self.time_total = t_wait_drone + self.t_off + self.t_exe
        else:
            self.time_total = self.t_local
        self.cost = self.time_total
        return self.cost


class Partition(Participant):

    def __init__(self, idx, pos, ues):
        super().__init__(idx, pos)
        self.ues = ues  # 当前 Part 包含的所有 UDs
        self.drone = 0  # 当前 Part 分配的 UAV ID
        self.t_wait = 0
        self.t_fly = 0

    def calc_cost(self, time_wait_drone, trans_rate, cpu_frequency):
        costs = 0
        for ue in self.ues:
            costs += ue.calc_cost(time_wait_drone, trans_rate[ue.idx], cpu_frequency[ue.idx])
        return costs

    def calc_exe_time(self, time_wait_drone, trans_rate, cpu_frequency):
        exe_times = [ue.calc_exe_time(time_wait_drone, trans_rate[ue.idx], cpu_frequency[ue.idx]) for ue in self.ues]
        return max(exe_times)


class Drone(Participant):

    def __init__(self, idx, pos, height, cpu_frequency, bandwidth, trans_power, flying_power, speed, battery):
        super().__init__(idx, pos)
        self.height = height                # 固定巡航高度，单位 M
        self.cpu_frequency = cpu_frequency  # CPU 频率，单位 MHz
        self.bandwidth = bandwidth          # 通信带宽，单位 MHz
        self.trans_power = trans_power      # 传输功率，单位 W
        self.flying_power = flying_power    # 平均飞行功率，计算飞行能耗
        self.speed = speed                  # 航行速度，单位 m/s
        self.battery = battery              # 初始电量，单位 mAh

    def update_pos(self, pos):
        self.pos[0] = pos[0]
        self.pos[1] = pos[1]
        return self.pos


class EnvConfig:

    def __init__(self, config_file):
        # 1. Task Config
        self.task_min_size = int(config_file['TASK']['min_size'])
        self.task_max_size = int(config_file['TASK']['max_size'])
        # 2. UEs Config
        self.ue_num = int(config_file['UE']['num'])
        self.ue_strategy = config_file['UE']['strategy']
        self.ue_cpu_frequency = int(config_file['UE']['cpu_frequency'])
        self.ue_trans_power = int(config_file['UE']['trans_power'])
        self.ue_coefficient_performance = float(config_file['UE']['coefficient_performance'])
        # 3. UAVs Config
        self.uav_num = int(config_file['UAV']['num'])
        self.uav_strategy = config_file['UAV']['strategy']
        self.uav_height = int(config_file['UAV']['height'])
        self.uav_radius = int(config_file['UAV']['radius'])
        self.uav_cpu_frequency = int(config_file['UAV']['cpu_frequency'])
        self.uav_bandwidth = int(config_file['UAV']['bandwidth'])
        self.uav_trans_power = int(config_file['UAV']['trans_power'])
        self.uav_flying_power = int(config_file['UAV']['flying_power'])
        self.uav_speed = int(config_file['UAV']['speed'])
        self.uav_battery = int(config_file['UAV']['battery'])
        self.uav_hover_time = int(config_file['UAV']['hover_time'])
        # 4. Env Config
        self.area = int(config_file['area'])
        self.area_std = (int(self.area / 10), int(self.area - self.area / 10))
        self.channel_gain = float(config_file['channel_gain'])
        self.channel_noise = float(config_file['channel_noise'])
        self.average_density = int(config_file['average_density'])
        self.episode = int(config_file['episode'])
        self.cd_episode = int(config_file['cd_episode'])


class Env:

    def __init__(self, config_file):
        # 1. 读取配置
        self.config = EnvConfig(config_file)
        self.M, self.N = self.config.ue_num, self.config.uav_num
        self.action, self.bandwidth, self.frequency, self.path = [], [], [], []     # 决策变量
        # 2. 初始化 UEs，按照一定聚集性随机初始化
        self.meta_ues, self.ues = self.init_ues()
        self._tasks = None
        # 3. 初始化 UAVs
        self.meta_drones, self.drones = self.init_drones()
        # 4. 初始化分区，使用迭代 + KMeans 实现几何覆盖
        self.G, self.meta_parts, self.parts, self.parts_ues = None, None, None, None
        # 5. 初始化 Parts 和 UAVs 的距离图和成本图（均为邻接矩阵）
        self.graph, self.cost_graph = None, None
        # 6. 初始化元数据
        self.meta_task_size, self.meta_task_cpu, self.meta_ue_loc, self.meta_trans_rate = None, None, None, None

    def init_ues(self):
        # 1. 随机生成 UEs 聚集的中心点
        average = np.ceil(self.M / self.config.average_density).astype(int)
        center_num = average + np.random.randint(1, average + 1)
        ues_pos, ues_features = make_blobs(
            n_samples=self.M, centers=center_num, center_box=self.config.area_std, random_state=0, cluster_std=100)
        # 2. 对超出边界的 UEs 位置进行规范化
        ues_pos = np.array([[max(10, min(ue_pos[0], self.config.area - 10)),
                             max(10, min(ue_pos[1], self.config.area - 10))] for ue_pos in ues_pos])
        ues_pos = np.concatenate((ues_pos, np.zeros(self.M)[:, np.newaxis]), axis=1)
        # 3. 初始化每个 UDs 及其计算任务
        ues = []
        for i in range(self.M):
            ues.append(UE(i, ues_pos[i], task=None,
                          cpu_frequency=self.config.ue_cpu_frequency,
                          trans_power=self.config.ue_trans_power,
                          coefficient_performance=self.config.ue_coefficient_performance))
        return ues_pos, np.array(ues)

    def init_drones(self):
        # 1. 随机生成 UAVs 的位置
        drones_pos = np.random.randint(self.config.area_std[0], self.config.area_std[1], size=(self.N, 2))
        drones_height = (np.ones(self.N) * self.config.uav_height)[:, np.newaxis]
        drones_pos = np.concatenate((drones_pos, drones_height), axis=1)
        # 2. 初始化每个 UAVs
        drones = []
        for i in range(self.N):
            drones.append(Drone(idx=i, pos=drones_pos[i],
                                height=self.config.uav_height,
                                cpu_frequency=self.config.uav_cpu_frequency,
                                bandwidth=self.config.uav_bandwidth,
                                trans_power=self.config.uav_trans_power,
                                flying_power=self.config.uav_flying_power,
                                speed=self.config.uav_speed,
                                battery=self.config.uav_battery))
        return drones_pos, np.array(drones)

    def generate_task(self):
        tasks = []
        for i in range(self.M):
            task = Task.factory(i, self.config.task_min_size, self.config.task_max_size)
            self.ues[i].set_task(task)
            tasks.append([task.size, task.cycles_per_bit, task.weight_time, task.weight_energy])
        return np.array(tasks)

    def partition(self, partition_strategy):
        print("********************************************************************************")
        print(f"Algorithm.1 Start partitions ...")
        now = util.now()
        # 1. 调用分区策略获取分区位置和分区结果
        circles_pos, partition_group = partition_strategy.run(self.meta_ues)
        # 2. 为 UDs 设置分区信息
        parts = []
        g = len(circles_pos)
        for part_idx in range(g):
            ues_idx = partition_group[part_idx]
            parts.append(Partition(part_idx, circles_pos[part_idx], self.ues[ues_idx]))
            for ue_id in ues_idx:
                self.ues[ue_id].part = part_idx
        print(f"End => g={g}, waste time: {util.now() - now}ms")
        # 3. 保存为环境信息
        self.G, self.meta_parts, self.parts, self.parts_ues = g, circles_pos, np.array(parts), partition_group
        # 4. 创建分区的成本邻接矩阵
        self.graph, self.cost_graph = self.create_graph(), []

    def new_task(self):
        self._tasks = self.generate_task()

    def execute(self, action=None, bandwidth=None, frequency=None):
        if action is not None:
            self.action = action
            for i in range(self.M):
                self.ues[i].offloading = self.action[i]
        if bandwidth is not None:
            self.bandwidth = bandwidth
        if frequency is not None:
            self.frequency = frequency

    def calc_path_cost(self, path=None, execute=False):
        # 1. 初始化总成本与当前等待时间
        if not path:
            path = self.path
        cost, time_wait_drone = 0, 0
        drone = self.drones[0]
        drone_pos = self.parts[0].pos
        # 2. 按照给定路径依次处理每个 PT
        for part_idx in path:
            # 2.1. 获取 PT 的信息，计算等待时间
            part = self.parts[part_idx]
            distance = np.linalg.norm(drone_pos - part.pos).astype(int)
            time_wait_drone += np.round(distance / self.config.uav_speed, decimals=2)
            part.time_wait_drone = time_wait_drone
            # 2.2. 更新 UAV 位置到 PT 中心
            drone_pos = part.pos
            drone.pos[0] = part.pos[0]
            drone.pos[1] = part.pos[1]
            # 2.3. 计算 PT 内所有 UDs 的
            time_service = 0
            for ue in part.ues:
                distance_ue = drone.calc_distance(ue)
                channel_gain = self.config.channel_gain / (distance_ue ** 2)
                trans_rate = self.bandwidth[ue.idx] * np.log2(
                    1 + (drone.trans_power * channel_gain) / self.config.channel_noise)
                cost += ue.calc_cost(time_wait_drone, trans_rate, self.frequency[ue.idx]) \
                    if not execute else ue.execute_action(time_wait_drone, trans_rate, self.frequency[ue.idx])
                time_service = max(time_service, ue.t_off + ue.t_exe)
            time_wait_drone += time_service
        return cost

    def calc_cost(self, action=None, bandwidth=None, frequency=None, path=None, execute=False):

        action = (self.action if action is None else action)
        bandwidth = (self.bandwidth if bandwidth is None else bandwidth)
        frequency = (self.frequency if frequency is None else frequency)
        path = (self.path if path is None else path)

        # 1. T_WAIT
        t_fly = np.zeros(self.G, dtype=int)
        for idx, g in enumerate(path):
            if idx == 0:
                continue
            distance = np.linalg.norm(self.meta_parts[path[idx-1]] - self.meta_parts[g]).astype(int)
            t_fly[idx] = np.round(distance / self.config.uav_speed, decimals=2)
        t_wait = np.zeros(self.G, dtype=int)
        for idx, g in enumerate(path):
            t_wait[g] = idx * self.config.uav_hover_time
            for j in range(idx + 1):
                t_wait[g] += t_fly[j]
        t_wait_m = np.zeros(self.M, dtype=int)
        for m in range(self.M):
            t_wait_m[m] = t_wait[self.ues[m].part]

        # 2. T_OFF, T_EXE
        with np.errstate(divide='ignore'):
            t_off = self.meta_trans_rate / bandwidth
            t_exe = self.meta_task_cpu / frequency
            for m in range(self.M):
                if action[m] == 0:
                    t_off[m] = 0
                    t_exe[m] = 0

        # 3. Cost
        cost = (action * (t_wait_m + t_off + t_exe) / self.meta_ue_loc) + (1 - action)
        return sum(cost)

    def calc_utilization(self):
        # 1. 计算边缘计算时间
        with np.errstate(divide='ignore'):
            t_off = self.meta_trans_rate / self.bandwidth
            t_exe = self.meta_task_cpu / self.frequency
            for m in range(self.M):
                if self.action[m] == 0:
                    t_off[m] = 0
                    t_exe[m] = 0
        t_edge = t_off + t_exe
        # 2. 统计每个分区的计算时间、浪费时间及飞行时间
        t_fly = np.zeros(self.G, dtype=float)
        t_use = np.zeros(self.G, dtype=float)
        for idx, g in enumerate(self.path):
            ues_idx = self.parts_ues[g]
            t_use[g] = np.max(t_edge[ues_idx])
            if idx == 0:
                continue
            distance = np.linalg.norm(self.meta_parts[self.path[idx - 1]] - self.meta_parts[g]).astype(int)
            t_fly[g] = np.round(distance / self.config.uav_speed, decimals=2)
        # 3. 计算时间利用率
        t_total = self.config.uav_hover_time * self.G + sum(t_fly)
        t_ratio = round(sum(t_use) / t_total, 3)
        # 4. 计算 CPU 利用率
        cpu_use = np.zeros(self.G, dtype=float)
        for idx, g in enumerate(self.path):
            ues_idx = self.parts_ues[g]
            cpu_use[g] = sum(t_edge[ues_idx] * self.frequency[ues_idx])
        cpu_ratio = sum(cpu_use) / (self.config.uav_cpu_frequency * t_total)
        return t_ratio, cpu_ratio

    def create_graph(self, normalization=False):
        graph = np.zeros((self.G, self.G))
        for i in range(self.G):
            for j in range(self.G):
                if i != j:
                    graph[i][j] = self.parts[i].calc_distance(self.parts[j])
        return (np.round(graph / np.mean(graph), decimals=2)) if normalization else np.round(graph, decimals=2)

    def path_format(self):
        path_str = ""
        for i, p in enumerate(self.path):
            ues_idx = self.parts_ues[p]
            a = np.array(self.action)
            off_num = sum(a[ues_idx]) if len(self.action) == self.M else 0
            path_str += f"{p} ({len(self.parts_ues[p])}, {off_num})"
            if i != (len(self.path) - 1):
                path_str += " -> "
        return path_str

    def plot_distribution(self):
        plt.rcParams['xtick.direction'] = 'in'
        plt.rcParams['ytick.direction'] = 'in'
        plt.rcParams['font.size'] = 18
        colors = ['#3b6291', '#943c39', '#779043', '#624c7c', '#388498', '#bf7334', '#675083', '#9c403d', '#7d9847', '#3b8ba1', '#c97937']
        fig, ax = plt.subplots(figsize=(8, 8))
        for i in range(self.G):
            # UEs
            part_ues_idx = self.parts_ues[i]
            ax.scatter(x=self.meta_ues[part_ues_idx, 0].squeeze(), y=self.meta_ues[part_ues_idx, 1].squeeze(), marker='o', c=colors[i%len(colors)], s=25)
            # Parts
            center = self.meta_parts[i]
            draw_circle = plt.Circle((center[0], center[1]), radius=self.config.uav_radius, facecolor="None", edgecolor="#333", linestyle=':', linewidth=2)
            ax.annotate('{}'.format(i), (center[0], center[1]), textcoords="offset points", xytext=(0, 10), ha='center')
            ax.add_artist(draw_circle)

        # line_x = []
        # line_y = []
        # for g in self.path:
        #     line_x.append(self.meta_parts[g][0])
        #     line_y.append(self.meta_parts[g][1])
        # line = plt.Line2D(line_x, line_y, color='black', marker='o', linestyle='-.', markerfacecolor='none', markersize=20, linewidth=3)
        # ax.add_line(line)

        ax.set_xlabel('X(m)')
        ax.set_ylabel('Y(m)')
        ax.axis([0, self.config.area, 0, self.config.area])
        ax.grid(True, color="#eee")
        plt.show()

    def init_meta(self):
        matrix = np.array([[ue.task.size, ue.task.cycles_per_bit, ue.t_local] for ue in self.ues])
        task_size = matrix[:, 0].squeeze()
        task_cpu = matrix[:, 1].squeeze()
        ue_loc = matrix[:, 2].squeeze()
        trans_rate = np.zeros(self.M)
        for i in range(self.M):
            ue = self.ues[i]
            part_uav_pos = np.array([self.parts[ue.part].pos[0], self.parts[ue.part].pos[1], self.config.uav_height])
            distance = np.linalg.norm(ue.pos - part_uav_pos).astype(int)
            channel_gain = self.config.channel_gain / (distance ** 2)
            trans_rate[i] = np.log2(1 + (ue.trans_power * channel_gain) / self.config.channel_noise)
        trans_rate = task_size / trans_rate
        task_cpu = task_size * task_cpu
        self.meta_task_size, self.meta_task_cpu, self.meta_ue_loc, self.meta_trans_rate = task_size, task_cpu, ue_loc, trans_rate

    def generate_path(self):
        # 1. TSP result
        # self.path = [1, 2, 0, 3]
        self.path = [1, 5, 0, 6, 2, 4, 3]
        # self.path = [1, 0, 3, 5, 4, 2]
        # self.path = [1, 0, 5, 2, 6, 4, 3]
        # self.path = [1, 7, 0, 6, 8, 5, 3, 4, 2]
        # self.path = [3, 5, 4, 8, 2, 9, 7, 11, 10, 1, 0, 6]
        # self.path = [0, 1, 8, 3, 9, 7, 5, 4, 6, 2]
        # self.path = [2, 6, 10, 4, 9, 8, 3, 0, 7, 1, 5]
        # self.path = [1, 7, 3, 8, 2, 11, 10, 9, 5, 0, 6, 4]
        # self.path = [2, 4, 8, 6, 5, 1, 10, 9, 0, 7, 3]
        # 2. Calculate parts t_fly, t_wait
        t_fly = 0
        for idx, part in enumerate(self.path):
            if idx != 0:
                distance = self.parts[part].calc_distance(self.parts[self.path[idx-1]])
                t_fly_part = (distance / self.config.uav_speed)
                t_fly += t_fly_part
                self.parts[part].t_fly = t_fly_part
            self.parts[part].t_wait = idx * self.config.uav_hover_time + t_fly
